
from ultralytics import YOLO


def main(args):
    model = YOLO(args.model)

    # Key			Value	Description
    # data			None	path to data file, i.e. coco128.yaml
    # imgsz			640		size of input images as integer
    # batch			16		number of images per batch (-1 for AutoBatch)
    # save_json		False	save results to JSON file
    # save_hybrid	False	save hybrid version of labels (labels + additional predictions)
    # conf			0.001	object confidence threshold for detection
    # iou			0.6		intersection over union (IoU) threshold for NMS
    # max_det		300		maximum number of detections per image
    # half			True	use half precision (FP16)
    # device		None	device to run on, i.e. cuda device=0/1/2/3 or device=cpu
    # dnn			False	use OpenCV DNN for ONNX inference
    # plots			False	save plots and images during train/val
    # rect			False	rectangular val with each batch collated for minimum padding
    # split			val		dataset split to use for validation, i.e. 'val', 'test' or 'train'
    kwargs = {
        "data": args.data,
        # "imgsz": 640,
        # "batch": -1,

        # "half": True,
        "dnn": True,
        "conf": 0.01,
        # "iou": 0.6,
        # "max_det": 300,

        "device": args.device,

        # 是否保存图表
        "plots": False,
        "save_json": True,
        # 包括原始标签和额外的预测结果
        # "save_hybrid": True,
    }
    metrics = model.val(**kwargs)
    # ultralytics.utils.metrics.DetMetrics
    print('metrics', metrics)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('model', type=str)
    parser.add_argument('data', type=str)
    parser.add_argument('--device', type=str, default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--verbose', action='store_true')
    args = parser.parse_args()
    main(args)
